from manimlib import *
import numpy as np
import sys, os

sys.path.append(os.path.dirname(__file__))
from MySpring import *


class Spring2DStrecth(Scene):
    """
    2D弹簧压缩
    """

    def setup(self):
        self.axes = Axes(
            # x轴的范围从-1到10，步长为1
            x_range=(-1.5, 1.5, 0.1),
            # y轴的范围从-2到2，步长为0.5y-axis ranges from -2 to 10 with a step size of 0.5
            y_range=(-0.5, 0.5, 0.1),
            # 坐标系将会伸缩来匹配指定的height和width
            height=6,
            width=10,
            # Axes由两个NumberLine组成，你可以通过axis_config来指定它们的样式
            axis_config={
                "stroke_color": GREY_A,
                "stroke_width": 2,
            },
            # 或者，你也可以像这样分别指定各个坐标轴的样式
            y_axis_config={
                "include_tip": False,
            },
        )
        # add_coordinate_labels方法的关键字参数可以传入DecimalNumber来指定它的样式
        self.axes.add_coordinate_labels(
            font_size=20,
            num_decimal_places=1,
        )
        self.add(self.axes)
        pass

    def construct(self):
        # ========================================================
        # 1. 坐标系

        # ========================================================
        # ========================================================
        # ========================================================

        # 创建弹簧
        sprint1_start = [0, 0.2, 0]
        sprint1_end = [1.3, 0.3, 0]
        spring1 = Spring2D(sprint1_start, sprint1_end, self.axes)

        # 绘制弹簧
        self.play(ShowCreation(spring1))

        # 弹簧压缩
        self.play(
            spring1.animate.stretch(
                0.8, 0, about_point=self.axes.c2p(sprint1_start[0], sprint1_start[1], sprint1_start[2])
            )
        )

        self.wait()


class Spring3DStrecth(Scene):
    """
    3D弹簧压缩
    """

    def setup(self):
        self.axes = ThreeDAxes()

    def construct(self):
        # ========================================================
        # 1. 坐标系

        # ========================================================
        # ========================================================
        # ========================================================

        # 创建弹簧
        sprint1_start = [0, 0, 0]
        sprint1_end = [1, 0, 0]
        spring1 = Spring3D(sprint1_start, sprint1_end, self.axes)

        # 绘制弹簧
        self.play(ShowCreation(spring1))

        # 弹簧压缩
        self.play(
            spring1.animate.stretch(
                0.8, 0, about_point=self.axes.c2p(sprint1_start[0], sprint1_start[1], sprint1_start[2])
            )
        )

        self.wait()
